/*
 *     Copyright (C) <2021>  <Junwen Chen>
 *
 *     This program is free software: you can redistribute it and/or modify
 *     it under the terms of the GNU General Public License as published by
 *     the Free Software Foundation, either version 3 of the License, or
 *     (at your option) any later version.
 *
 *     This program is distributed in the hope that it will be useful,
 *     but WITHOUT ANY WARRANTY; without even the implied warranty of
 *     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *     GNU General Public License for more details.
 *
 *     You should have received a copy of the GNU General Public License
 *     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

package io.ordinate.engine.physicalplan;

import io.ordinate.engine.record.RootContext;
import io.reactivex.rxjava3.core.Observable;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.Schema;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collections;
import java.util.List;
import java.util.stream.IntStream;

public class OutputLinq4jPhysicalPlan implements PhysicalPlan {
    private static final Logger LOGGER = LoggerFactory.getLogger(OutputLinq4jPhysicalPlan.class);

    public static OutputLinq4jPhysicalPlan create(PhysicalPlan input) {
        return new OutputLinq4jPhysicalPlan(input);
    }

    public OutputLinq4jPhysicalPlan(PhysicalPlan input) {
        this.input = input;
    }

    final PhysicalPlan input;

    @Override
    public Schema schema() {
        return this.input.schema();
    }

    @Override
    public List<PhysicalPlan> children() {
        return Collections.singletonList(input);
    }

    @Override
    public Observable<VectorSchemaRoot> execute(RootContext rootContext) {
        return this.input.execute(rootContext);
    }

    @Override
    public void accept(PhysicalPlanVisitor physicalPlanVisitor) {
        physicalPlanVisitor.visit(this);
    }

    public Observable<Object[]> executeToObject(RootContext rootContext) {
        return input.execute(rootContext).flatMap(c -> {
            int rowCount = c.getRowCount();
            int columnCount = c.getFieldVectors().size();
            Object[][] array = IntStream.range(0, rowCount).mapToObj(i -> new Object[columnCount]).toArray(n -> new Object[rowCount][]);
            for (int columnId = 0; columnId < columnCount; columnId++) {
                FieldVector vector = c.getVector(columnId);
                for (int rowId = 0; rowId < rowCount; rowId++) {
                    Object object = vector.getObject(rowId);
                    array[rowId][columnId] = object;
                }
            }
            return Observable.fromArray(array);
        });
    }
}
